import os
import numpy as np
import pandas as pd
import pywt

def extract_features(signal, fs, start_index):
    features = {}

    # 时间（分钟）计算：每个点为 1 / (fs × 60) 分钟，加上起始偏移 74 分钟
    features['time_min'] = start_index / (fs * 60) + 74

    # 小波分解
    coeffs = pywt.wavedec(signal, 'db6', level=5)
    cA5, cD5, cD4, cD3, cD2, cD1 = coeffs

    def energy(x):
        return np.sum(x ** 2)

    # 特征提取
    features['level1d_abs_mean'] = np.mean(np.abs(cD1))
    features['band_energy_30_35Hz'] = energy(signal[int(30 * len(signal) / fs):int(35 * len(signal) / fs)])
    features['band_energy_65_70Hz'] = energy(signal[int(65 * len(signal) / fs):int(70 * len(signal) / fs)])
    features['level2d_abs_mean'] = np.mean(np.abs(cD2))
    features['abs_mean'] = np.mean(np.abs(signal))
    features['mean_amp'] = np.mean(signal)
    features['level1d_energy'] = energy(cD1)
    features['rms'] = np.sqrt(np.mean(signal ** 2))
    features['std'] = np.std(signal)
    features['frf'] = features['rms']  # 方根幅值
    features['level5d_abs_mean'] = np.mean(np.abs(cD5))
    features['level2d_energy'] = energy(cD2)
    features['level1d_energy_std'] = np.std(cD1 ** 2)
    features['waveform_factor'] = features['rms'] / features['abs_mean'] if features['abs_mean'] != 0 else 0

    return features

def sliding_window_features(signal, fs, window_size=2048, step_size=1024):
    feature_rows = []
    for start in range(0, len(signal) - window_size + 1, step_size):
        window = signal[start:start + window_size]
        feats = extract_features(window, fs, start)
        feature_rows.append(feats)
    return feature_rows

def process_multiple_files(input_dir, output_path, file_range, sampling_freq):
    combined_signal = []

    # 拼接所有第75到123个CSV的数据
    for i in range(file_range[0], file_range[1] + 1):
        if i < 75:
            continue  # 跳过前74个文件

        file_name = f"{i}.csv"
        file_path = os.path.join(input_dir, file_name)

        if not os.path.isfile(file_path):
            print(f"文件未找到：{file_name}，跳过")
            continue

        data = pd.read_csv(file_path, header=None, skiprows=1)
        data[0] = pd.to_numeric(data[0], errors='coerce')
        data = data.dropna(subset=[0])
        signal = data.iloc[:, 0].values.astype(float)
        combined_signal.extend(signal)

        print(f"已加载文件 {i}: {file_name}, 样本数: {len(signal)}")

    print(f"\n拼接后总长度: {len(combined_signal)} 点")

    # 滑动窗口提取特征
    all_features = sliding_window_features(
        signal=np.array(combined_signal),
        fs=sampling_freq,
        window_size=2048,
        step_size=1024
    )

    df_all = pd.DataFrame(all_features)
    df_all.to_csv(output_path, index=False)
    print(f"\n所有特征已保存到：{output_path}")


input_directory = r"D:\轴承寿命预测\Bearing1_1"
output_file = "./combined_features_fixed.csv"

process_multiple_files(
    input_dir=input_directory,
    output_path=output_file,
    file_range=(1, 123),     # 文件范围
    sampling_freq=25600      # 采样频率
)